import torch
# from dig.threedgraph.dataset import QM93D
from data_processing import MD17
from model.method import SphereNet #SchNet, DimeNetPP, ComENet
from ..aimadv import adversarial_train
from ..aimadv.eval import ThreeDEvaluator

device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device("cpu")

dataset_md17 = MD17(root='dataset/', name='aspirin')
print(dataset_md17.data)

split_idx_md17 = dataset_md17.get_idx_split(len(dataset_md17.data.y), train_size=1000, valid_size=1000, seed=42)

train_dataset_md17, valid_dataset_md17, test_dataset_md17 = dataset_md17[split_idx_md17['train']], dataset_md17[split_idx_md17['valid']], dataset_md17[split_idx_md17['test']]
print('total, train, validaion, test:',len(dataset_md17.data.y), len(train_dataset_md17), len(valid_dataset_md17), len(test_dataset_md17))

model_md17 = SphereNet(energy_and_force=True, cutoff=5.0, num_layers=4, 
        hidden_channels=128, out_channels=1, int_emb_size=64, 
        basis_emb_size_dist=8, basis_emb_size_angle=8, basis_emb_size_torsion=8, out_emb_channels=256, 
        num_spherical=3, num_radial=6, envelope_exponent=5, 
        num_before_skip=1, num_after_skip=2, num_output_layers=3 
        )
loss_func_md17 = torch.nn.L1Loss()
evaluation_md17 = ThreeDEvaluator()
 
run3d_md17 = adversarial_train()
run3d_md17.run(device, 
               train_dataset_md17, 
               valid_dataset_md17, 
               test_dataset_md17, 
               model_md17, 
               loss_func_md17, 
               evaluation_md17, 
               epochs=100, 
               batch_size=1, 
               vt_batch_size=64, 
               lr=0.0005, 
               lr_decay_factor=0.5, 
               lr_decay_step_size=200, 
               energy_and_force=True,
               save_dir="./result/checkpoint",
               log_dir="./result/logs")

